import asyncio
from typing import Dict, Literal, Any

from pydantic import BaseModel, Field


class ChatEvent(BaseModel):
    # 事件类型："ready", "update_session", "token", "complete", "error", "close"
    event_type: Literal["ready", "update_session", "tool_call", "tool_result", "token", "complete", "error", "close"]
    # 数据
    data: Dict[str, Any] = Field(default_factory=dict)
    # 事件戳
    timestamp: float = Field(default_factory=lambda: asyncio.get_event_loop().time())

    def formate_response(self):
        """
        格式化response
        :return: str
        """
        return f"event: {self.event_type}\ndata:{self.data} \n\n"

    def format_data(self):
        """
        data格式化
        :return: str
        """
        if self.event_type == "token":
            return self.data['token']
        elif self.event_type == "error":
            return self.data['error_msg']
        return None

